[CPU][Spec] Enable DFLASH speculative decoding on Xeon (intel_amx) - #36782

Motivation

This PR enables DFlash speculative decoding on Xeon CPUs which currently is rejected outright.

I could not find an existing PR for this. The closest related work is #35629
(NPU DFlash2, open), #27854 (AMD DFlash fused KV, merged) and #30964 (AMD
DSPARK, merged).

Modifications

The reason dflash_head_logits() is needed: on CPU + AMX, ParallelLMHead gets
PackWeightMethod, so lm_head.weight is VNNI-packed — but
should_apply_lm_head_quant_method is False and is_dense_head_weight is
True, so the existing dense matmul paths read packed bytes as if they were
plain bf16 and produce garbage. VocabParallelEmbedding is not packed, so the
embedding path is unaffected.

Two new registered CPU unit tests.
test/registered/unit/spec/test_dflash_cpu.py covers the device gate, the CPU
overlap-schedule resolution, draft-backend resolution, the linear-chain
tree_topk, and the AMX-prepacked weight guards.
test/registered/unit/layers/attention/test_intel_amx_non_causal_mask.py pins
the new all-visible mask for an ENCODER_ONLY verify layer and, as a negative
control, that a causal layer still takes the mask-free kernel path.

Only intel_amx is enabled as a CPU draft backend. torch_native is
deliberately excluded: it is not validated here, and #36361 proposes rejecting
that backend with speculative decoding outright.

Accuracy Tests

Intel Xeon 6 (128 physical cores across 4 NUMA nodes),
meta-llama/Llama-3.1-8B-Instruct +
z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat:

Single rank

python -m sglang.launch_server \
  --model-path meta-llama/Llama-3.1-8B-Instruct \
  --speculative-algorithm DFLASH \
  --speculative-draft-model-path z-lab/LLaMA3.1-8B-Instruct-DFlash-UltraChat \
  --device cpu --attention-backend intel_amx --trust-remote-code \
  --disable-overlap-schedule --context-length 8192 \
  --max-running-requests 8 --mem-fraction-static 0.3

Initialized DFLASH draft runner. attention_backend=intel_amx,
model=DFlashDraftModel, block_size=10, draft_window_size=None, compact_cache=False
DFLASH draft runner ready. mask_token=<|MASK|>, mask_token_id=128002

At temperature: 0 the output was byte-identical with and without DFlash.

Speculative metrics for that run:

spec_accept_rate 0.185, spec_accept_length 2.667, spec_verify_ct 12,
spec_num_correct_drafts 20, spec_num_proposed_drafts 108,
spec_correct_drafts_histogram [3, 5, 0, 1, 3]

Tensor parallel

Repeated at --tp 4 with SGLANG_CPU_OMP_THREADS_BIND="0-31|32-63|64-95|96-127"
(one rank per NUMA node), since CPU CI runs TP > 1. The greedy output matches the
non-speculative run token for token, so the sharded ParallelLMHead and
qkv_proj paths behave. The one class of exception is characterised below.

GSM8K, 200 questions, --max-new-tokens 128:

Each delta is 0.005, a single question, in the direction of DFLASH.

Speed Tests and Profiling

HumanEval (first 20 problems), 128 output tokens with EOS ignored, greedy,
--max-concurrency 1, 3 warmup requests, --tp 4,
--mem-fraction-static 0.8, --speculative-dflash-block-size 9 (8 draft
tokens). Both configurations emit exactly 2560 tokens, so length cannot distort
the comparison.

python -m sglang.benchmark.serving --backend sglang \
  --model meta-llama/Llama-3.1-8B-Instruct \
  --dataset-name custom --dataset-path humaneval20.jsonl \
  --sharegpt-output-len 128 --num-prompts 20 --max-concurrency 1 \
  --request-rate inf --warmup-requests 3 --temperature 0 --seed 42

Mean of 3 runs per configuration:

Run-to-run spread is 1.2% on the speculative side and 0.3% on the baseline.

The result decomposes cleanly. A verify step costs
TPOT x accept_length = 12.23 x 3.75 = 45.9 ms against a 28.56 ms decode step,
a ratio of 1.61, so the predicted speedup is 3.75 / 1.61 = 2.33x against
2.33x measured. End-to-end speed is fully accounted for by acceptance length
and per-step overhead, with no unexplained residual.

Known limitations

• Hybrid linear-attention models (e.g. Qwen3.5) still cannot run DFlash on CPU. (Update: Currently planned via [CPU][Spec] Support DFlash target-verify for hybrid GDN models #37851 )
  The CPU GDN kernels have no target-verify entry points: causal_conv1d_update_cpu
  has no intermediate_conv_window argument and asserts x.dim() == 2,
  fused_sigmoid_gating_delta_rule_update_cpu has no disable_state_update /
  intermediate-state arguments, and scatter_mamba_states_after_mtp_verify is
  Triton-only. This blocks EAGLE and MTP on those models too, not just DFlash,
  so I plan to send it as a separate follow-up PR.
• The DFlash draft/target context-length check requires --context-length 8192
  (or SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1) for this checkpoint pair.
  That is pre-existing and reproduces on CUDA.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ❌ Run #33142487182
Latest PR Test (Extra): ❌ Run #33142486930
Latest PR Test (AMD ROCm 7.2): ❌ Run #33142487031